# coding=utf8
import sys

sys.path.append('../')
sys.path.append('../../')
import traceback
import copy
import time
from multiprocessing.dummy import Pool as ThreadPool  # 多线程

from i_entity_extractor.common_parser_lib.mongo import MongDb

mongo_conf37 = {
    'host': '101.201.102.37',
    'port': 28019,
    'final_db': 'final_data',
    'username': None,
    'password': None,
}


class CollectCompany:
    def __init__(self):
        self.db = MongDb(mongo_conf37['host'], mongo_conf37['port'], mongo_conf37['final_db'], mongo_conf37['username'],
                         mongo_conf37['password'])
        self.src_table = 'enterprise_data_gov'
        self.target_table = 'gsxx_company'

    def get_data(self):
        result = self.db.select(self.src_table, {})

        num = 0
        company_list = []
        begin_time = time.time()
        pool = ThreadPool(4)
        for item in result:
            company_map = {}
            try:
                company = item.get('company', '')
                if not company:
                    continue
                num += 1
                branch = item.get('branch', [])
                shareholder_information = item.get('branch', [])
                invested_companies = item.get('invested_companies', [])

                company_map['relate_company'] = []

                for info in branch:
                    company_name = info.get('company_name', '')
                    if company_name:
                        company_map['relate_company'].append(company_name)
                for info in shareholder_information:
                    company_name = info.get('shareholder_name', '')
                    if company_name:
                        company_map['relate_company'].append(company_name)
                for info in invested_companies:
                    company_name = info.get('name', '')
                    if company_name:
                        company_map['relate_company'].append(company_name)

                company_map['registered_capital'] = item.get('registered_capital', -1)
                company_map['province'] = item.get('province', '')
                company_map['registered_code'] = item.get('registered_code', '')
                company_map['company'] = company
                company_map['_id'] = company
                company_list.append(company_map)

                if num % 100 == 0:
                    results = pool.map(self.insert_mongo, company_list)

                    company_list = []
                    print num, "time_cost:", time.time() - begin_time
                if num >= 200:
                    break
            except:
                print traceback.format_exc()
        pool.close()
        pool.join()

    def insert_mongo(self, data_map):
        '''插入mongo'''
        self.db.upsert(self.target_table, data_map)


if __name__ == "__main__":
    obj = CollectCompany()
    obj.get_data()
